"""
IoU (Intersection over Union) / Jaccard 指标计算
"""
import torch
import numpy as np


def calculate_iou(pred, target, threshold=0.5, smooth=1e-6):
    """
    计算单个样本的IoU

    Args:
        pred: 预测结果 (H, W) 或 (1, H, W)，值在[0, 1]之间
        target: 真实标签 (H, W) 或 (1, H, W)，值为0或1
        threshold: 二值化阈值
        smooth: 平滑项，避免除零

    Returns:
        iou: IoU值
    """
    # 确保是numpy数组
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()

    # 展平
    pred = pred.flatten()
    target = target.flatten()

    # 二值化预测
    pred = (pred > threshold).astype(np.float32)

    # 计算交集和并集
    intersection = np.sum(pred * target)
    union = np.sum(pred) + np.sum(target) - intersection

    # 计算IoU
    iou = (intersection + smooth) / (union + smooth)

    return iou


def calculate_miou(pred, target, num_classes=2, threshold=0.5, smooth=1e-6):
    """
    计算批次的平均IoU (mIoU)

    Args:
        pred: 预测结果 (B, 1, H, W) 或 (B, C, H, W)
        target: 真实标签 (B, 1, H, W)
        num_classes: 类别数（对于二分类，通常是2）
        threshold: 二值化阈值
        smooth: 平滑项

    Returns:
        miou: 平均IoU
        iou_per_class: 每个类别的IoU
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()

    batch_size = pred.shape[0]

    # 如果是多通道输出，取前景通道
    if pred.shape[1] > 1:
        pred = pred[:, 1:2, :, :]  # 取前景类

    # 二值化
    pred = (pred > threshold).astype(np.float32)

    # 计算每个类别的IoU
    iou_list = []

    for cls in range(num_classes):
        if cls == 0:
            # 背景类
            pred_cls = 1 - pred
            target_cls = 1 - target
        else:
            # 前景类
            pred_cls = pred
            target_cls = target

        intersection = np.sum(pred_cls * target_cls)
        union = np.sum(pred_cls) + np.sum(target_cls) - intersection

        iou = (intersection + smooth) / (union + smooth)
        iou_list.append(iou)

    miou = np.mean(iou_list)

    return miou, iou_list
